from PIL import Image
import numpy as np
import tensorflow as tf
from Core.Model import AlexNet
from Core import utils
import os
import pandas as pd

model_save_path = 'checkpoint/AlexNet.ckpt'
model = AlexNet()
model.load_weights(model_save_path)

preNum = int(input("input the number of test pictures:"))

# 检查是否有table表
if not os.path.exists('D:/SystemSoft/GitRepos/cnn-alex-net/train&predict/checkpoint/table.xlsx'):
    utils.generate_table()

excel = pd.read_excel('D:/SystemSoft/GitRepos/cnn-alex-net/train&predict/checkpoint/table.xlsx')
for i in range(preNum):
    image_path = input('input the path of test picture:')
    img = Image.open(image_path)
    img = img.resize((227, 227), Image.ANTIALIAS)
    # 归一化处理
    img_arr = np.array(img)
    img_arr = img_arr / 255.0
    # print(img_arr)
    img_arr = np.reshape(img_arr, (1, 227, 227, 3))
    result = model.predict(img_arr)
    pred = tf.argmax(result, axis=1)
    print('\n')
    pred_ = excel.loc[pred, 'labels']
    print(pred_)
